# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging

import tabulate
from tornado.ioloop import IOLoop

import coin.flow.subscriber as flow_subscriber
from coin.base.param_util import to_list
from coin.flow.simple_queue import SimpleQueue
from coin.exchange.bitflyer_v1.kr_rest.currency import BitflyerCurrency
from coin.proto.coin_query_pb2 import CurrencyBalance
from coin.strategy.mm.bitflyer_v1.order_subsystem import BitflyerOrderSubsystem
from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.strategy.mm.base.bitmex_strategy_base import (BitmexPositionBalancePrinter)
from coin.strategy.mm.subscription import SubscriptionRequest


def calculate_unrealized_pnl(pos, avg_price, mid_price):
  if pos == 0:
    return 0
  return pos * (mid_price - avg_price)


class BitflyerMarginPositionBalancePrinter(BitmexPositionBalancePrinter):
  def __init__(self,
               currency,
               products,
               bookmap,
               pos_bias_map=None,
               strategy_logger=None,
               logger=None):
    assert currency == BitflyerCurrency.FromStrCurrency('JPY')
    super().__init__(products,
                     bookmap,
                     logger=(logger or logging.getLogger('BitflyerPositionBalancePrinter')))
    self._currency = currency
    self._pos_bias_map = pos_bias_map or {}
    self._strategy_logger = strategy_logger

  def _get_pos(self, product):
    book = self._bookmap.get_book(product)
    if book is None or not book.has_ask() or not book.has_bid():
      return None, None

    pos = self._og.get_position(product)
    average_entry_price = self._og.get_average_entry_price(product)

    midp = (book.ask0().price + book.bid0().price) / 2.
    unrealized_pnl = calculate_unrealized_pnl(pos, average_entry_price, midp)

    if product in self._pos_bias_map and self._pos_bias_map[product] != 0:
      pos_bias, _ = self._pos_bias_map[product]
      unrealized_pnl_from_bias = calculate_unrealized_pnl(pos_bias, 0., midp)
      return pos + pos_bias, unrealized_pnl + unrealized_pnl_from_bias
    else:
      return pos, unrealized_pnl

    return pos, unrealized_pnl

  def print_position_balance(self):
    if self._og is None or not self._og.is_ready():
      return

    table = []
    total_upnl = 0.
    for product in self._products:
      pos, unrealized_pnl = self._get_pos(product)
      total_upnl += (unrealized_pnl or 0.)
      table.append(['*', str(product), pos, unrealized_pnl])

    balance = self._og.get_balance(self._currency)

    table.append(['*', self._currency, None, balance])
    table.append(['*', 'Total', None, total_upnl + balance])

    pos_bal_table_str = tabulate.tabulate(table,
                                          headers=['*', 'Product', 'Position', 'Value'],
                                          floatfmt='.3f')
    self._logger.info('----- Position and balance:\n* %s\n%s',
                      str(datetime.datetime.now()),
                      pos_bal_table_str)

    if self._strategy_logger:
      pnl_balance = CurrencyBalance(currency='JPY', total=total_upnl + balance)
      self._strategy_logger.write_pnl_balance(pnl_balance=pnl_balance, quote=None, subreq_id=None)


class BitflyerStrategyBase(object):
  OGName = SubscriptionRequest('Futures', 'Bitflyer', 'v1')

  def __init__(self, products, currencies, logger=None):
    self._logger = logger or logging.getLogger('BitflyerStrategyBase')
    self._products = products
    self._currencies = sorted(to_list(currencies))
    self._ioloop = None
    self._queue = None
    self._flow_sub = None

    self._feed_subsystem = None
    self._og = None

  @property
  def order_gateway(self):
    return self._og

  @property
  def feed_subsystem(self):
    return self._feed_subsystem

  def register_flow(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._queue = queue
    self._flow_sub = flow_sub

  def register_feed_subsystem(self, feed_subsystem):
    self._feed_subsystem = feed_subsystem

  def register_order_gateway(self, name, order_gateway):
    self._logger.info('OG initialized: %s', name)
    if name == self.OGName:
      self._og = order_gateway
    else:
      raise ValueError('Unknown OG: %s', name)

    self._ioloop.add_callback(self.on_og_reset)

  def prepare(self):
    return

  def on_og_reset(self):
    raise NotImplementedError()

  def is_og_ready(self):
    return (self._og is not None and self._og.is_ready())


def run_live(strategy,
             og_config_name,
             init_order_gateway_after=None,
             exit_after=None,
             fs_config=None):
  # Initialize loop
  ioloop = IOLoop.current()
  queue = SimpleQueue('queue')
  flow_sub = flow_subscriber.from_queue(queue)
  strategy.register_flow(ioloop, queue, flow_sub)

  # Initialize feed
  feed_subsystem = MultiFeedSubsystem(fs_config=fs_config)
  strategy.register_feed_subsystem(feed_subsystem)

  # Initialize OG live environment
  env = BitflyerOrderSubsystem(currencies=strategy._currencies,
                               trade_products=strategy._products,
                               og_config_name=og_config_name,
                               init_og_after=init_order_gateway_after,
                               ioloop=ioloop,
                               queue=queue,
                               flow_sub=flow_sub,
                               og_log_callback=strategy.strategy_logger.write_og_log)
  env.on_order_gateway_reset_callback.append(strategy.register_order_gateway)

  # Start
  strategy.prepare()
  feed_subsystem.start(ioloop, queue, flow_sub)
  env.start()

  if exit_after is not None:
    ioloop.add_timeout(exit_after, ioloop.stop)

  try:
    ioloop.start()
    return 0
  except KeyboardInterrupt:
    return 1
